import os
import cv2
import math
import numpy as np
import torch
from pytorch_msssim import ssim

dir1 = '/root/code/FeMasr/result/NH-HAZE/restored/'
dir2 = '/root/data/NH-HAZE/test/clean/'
def compute_psnr(img1, img2):
    mse = np.mean((img1 - img2) ** 2)
    if mse == 0:
        return float('inf')
    return 20 * math.log10(255.0) - 10 * math.log10(mse)

def load_image(path):
    img = cv2.imread(path)
    if img is None:
        raise ValueError(f"Image {path} cannot be read")
    # Convert from BGR to RGB
    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

files1 = sorted(os.listdir(dir1))
files2 = sorted(os.listdir(dir2))
common_files = sorted(list(set(files1) & set(files2)))

if not common_files:
    print("No matching files found.")
    exit()

total_psnr = 0
total_ssim = 0
count = 0

for file in common_files:
    path1 = os.path.join(dir1, file)
    path2 = os.path.join(dir2, file)
    
    img1 = load_image(path1)
    img2 = load_image(path2)
    
    if img1.shape != img2.shape:
        print(f"Skipping {file}: shape mismatch")
        continue
    
    psnr_val = compute_psnr(img1, img2)
    
    # Prepare images for ssim computation
    img1_t = torch.from_numpy(img1).permute(2, 0, 1).unsqueeze(0).float() / 255.0
    img2_t = torch.from_numpy(img2).permute(2, 0, 1).unsqueeze(0).float() / 255.0
    
    ssim_val = ssim(img1_t, img2_t, data_range=1.0).item()
    
    print(f"{file}: PSNR = {psnr_val:.2f}, SSIM = {ssim_val:.4f}")
    
    if psnr_val != float('inf'):
        total_psnr += psnr_val
    total_ssim += ssim_val
    count += 1

if count > 0:
    avg_psnr = total_psnr / count
    avg_ssim = total_ssim / count
    print(f"Average PSNR: {avg_psnr:.2f}")
    print(f"Average SSIM: {avg_ssim:.4f}")
else:
    print("No images processed")